#include "Init.h"

#include "AttachDetachFilter.h"
#include "FastIO.h"
#include "Dispatch.h"
#include "Rules.h"
#include "Utils.h"

#include "Interface.h"

// global variables
PDRIVER_OBJECT g_pFilterDriverObject;
PDEVICE_OBJECT g_pFilterCDO;
FAST_MUTEX g_FilterCDOLock;

// add DriverEntry to INIT section: free memory of DriverEntry after call
#pragma alloc_text(INIT, DriverEntry)

NTSTATUS DriverEntry(IN PDRIVER_OBJECT pDriverObject, IN PUNICODE_STRING pRegistryPath)
{
    NTSTATUS ntStatus = STATUS_UNSUCCESSFUL;
    UNICODE_STRING usCDOName = {0};
    unsigned int uiMJCounter = 0;

    RtlInitUnicodeString(&usCDOName, FileHidingFilterCDOName);

    ntStatus = IoCreateDevice(pDriverObject,
                              0,
                              &usCDOName,
                              FILE_DEVICE_FILE_HIDING_FILTER,
                              FILE_DEVICE_SECURE_OPEN,
                              FALSE,
                              &g_pFilterCDO);

    if (!NT_SUCCESS(ntStatus))
    {
        return ntStatus;
    }

    g_pFilterDriverObject = pDriverObject;

    for (uiMJCounter = 0; uiMJCounter <= IRP_MJ_MAXIMUM_FUNCTION; ++uiMJCounter)
    {
        pDriverObject->MajorFunction[uiMJCounter] = FilterDispatchPassThrough;
    }

    pDriverObject->FastIoDispatch = ExAllocatePoolWithTag(NonPagedPool, sizeof(FAST_IO_DISPATCH), FileHidingFilterPoolTag);

    if (NULL == pDriverObject->FastIoDispatch)
    {
        IoDeleteDevice(g_pFilterCDO);

        return STATUS_INSUFFICIENT_RESOURCES;
    }

    RtlZeroMemory(pDriverObject->FastIoDispatch, sizeof(FAST_IO_DISPATCH));

    pDriverObject->FastIoDispatch->SizeOfFastIoDispatch = sizeof(FAST_IO_DISPATCH);
    pDriverObject->FastIoDispatch->FastIoCheckIfPossible = FastIOCheckIfPossible;
    pDriverObject->FastIoDispatch->FastIoRead = FastIORead;
    pDriverObject->FastIoDispatch->FastIoWrite = FastIOWrite;
    pDriverObject->FastIoDispatch->FastIoQueryBasicInfo = FastIOQueryBasicInfo;
    pDriverObject->FastIoDispatch->FastIoQueryStandardInfo = FastIOQueryStandardInfo;
    pDriverObject->FastIoDispatch->FastIoLock = FastIOLock;
    pDriverObject->FastIoDispatch->FastIoUnlockSingle = FastIOUnlockSingle;
    pDriverObject->FastIoDispatch->FastIoUnlockAll = FastIOUnlockAll;
    pDriverObject->FastIoDispatch->FastIoUnlockAllByKey = FastIOUnlockAllByKey;
    pDriverObject->FastIoDispatch->FastIoDeviceControl = FastIODeviceControl;
    pDriverObject->FastIoDispatch->FastIoDetachDevice = FastIODetachDevice;
    pDriverObject->FastIoDispatch->FastIoQueryNetworkOpenInfo = FastIOQueryNetworkOpenInfo;
    pDriverObject->FastIoDispatch->MdlRead = MdlRead;
    pDriverObject->FastIoDispatch->MdlReadComplete = MdlReadComplete;
    pDriverObject->FastIoDispatch->PrepareMdlWrite = PrepareMdlWrite;
    pDriverObject->FastIoDispatch->MdlWriteComplete = MdlWriteComplete ;
    pDriverObject->FastIoDispatch->FastIoReadCompressed = FastIoReadCompressed;
    pDriverObject->FastIoDispatch->FastIoWriteCompressed = FastIoWriteCompressed;
    pDriverObject->FastIoDispatch->MdlReadCompleteCompressed = MdlReadCompleteCompressed;
    pDriverObject->FastIoDispatch->MdlWriteCompleteCompressed = MdlWriteCompleteCompressed;
    pDriverObject->FastIoDispatch->FastIoQueryOpen = FastIoQueryOpen;

    InitializeRulesManager(pRegistryPath);
    ReloadRules();

    // initialize global CDO lock
    ExInitializeFastMutex(&g_FilterCDOLock);

    // register callback for FS notifications
    ntStatus = IoRegisterFsRegistrationChange(pDriverObject, FSNotificationCallback);

    if (!NT_SUCCESS(ntStatus))
    {
        IoDeleteDevice(g_pFilterCDO);
        
        return ntStatus;
    }

    return ntStatus;
}

#pragma alloc_text(PAGE, FSNotificationCallback)

void FSNotificationCallback(IN PDEVICE_OBJECT pTargetCDO, IN BOOLEAN bFSActive)
{
    PAGED_CODE();

    if (bFSActive)
    {
        NTSTATUS ntStatus = AttachToFSDeviceObject(pTargetCDO);

        ASSERT(NT_SUCCESS(ntStatus));

        ntStatus = EnumerateFSVolumesAndAttach(pTargetCDO);

        ASSERT(NT_SUCCESS(ntStatus));
    }
    else
    {
        DetachFilter(pTargetCDO);
    }
}

NTSTATUS FilterDispatchPassThrough(IN PDEVICE_OBJECT pDeviceObject, IN PIRP pIRP)
{
    PFILE_HIDING_FILTER_DEVICE_EXTENSION pFilterDeviceExt = (PFILE_HIDING_FILTER_DEVICE_EXTENSION)pDeviceObject->DeviceExtension;
    PIO_STACK_LOCATION pIOStack = IoGetCurrentIrpStackLocation(pIRP);

    if (pDeviceObject == g_pFilterCDO)
    {
        pIRP->IoStatus.Information = 0;
        pIRP->IoStatus.Status = STATUS_INVALID_DEVICE_REQUEST;

        switch (pIOStack->MajorFunction)
        {
            case IRP_MJ_CREATE:

            case IRP_MJ_CLOSE:

            case IRP_MJ_CLEANUP:
            {
                pIRP->IoStatus.Information = 0;
                pIRP->IoStatus.Status = STATUS_SUCCESS;

                break;
            }

            case IRP_MJ_DEVICE_CONTROL:
            {
                ULONG ulControlCode = pIOStack->Parameters.DeviceIoControl.IoControlCode;

                // dispatch user I/O control
                switch(ulControlCode)
                {
                    case IOCTL_FHF_RELOAD_RULES:
                    {
                        ReloadRules();

                        pIRP->IoStatus.Information = 0;
                        pIRP->IoStatus.Status = STATUS_SUCCESS;

                        break;
                    }
                 }
            }
        }

        return CompleteIRP(pIRP, IO_NO_INCREMENT);
    }

    if (NULL == pFilterDeviceExt)
    {
        return SetStatusAndCompleteIRP(pIRP, STATUS_INVALID_DEVICE_REQUEST, 0, IO_NO_INCREMENT);
    }

    switch (pIOStack->MajorFunction)
    {
        case IRP_MJ_CREATE:
        {
            return pFilterDeviceExt->Dispatch.pDispatchCreate(pDeviceObject, pIRP);
        }
        case IRP_MJ_CLOSE:
        {
            return pFilterDeviceExt->Dispatch.pDispatchClose(pDeviceObject, pIRP);
        }
        case IRP_MJ_DIRECTORY_CONTROL:
        {
            return pFilterDeviceExt->Dispatch.pDispatchDirectoryControl(pDeviceObject, pIRP);
        }
        default:
        {
            break;
        }
    }

    return PassThroughIRP(pIRP, pFilterDeviceExt->pNextLowerDevice);
}

PDEVICE_OBJECT GetFilterCDOAndLock()
{
    ExAcquireFastMutex(&g_FilterCDOLock);

    return g_pFilterCDO;
}

void ReleaseFilterCDOAndUnlock(IN OUT PDEVICE_OBJECT* ppFilterCDO)
{
    *ppFilterCDO = NULL;

    ExReleaseFastMutex(&g_FilterCDOLock);
}

PDRIVER_OBJECT GetFilterDriverObject()
{
    return g_pFilterDriverObject;
}

#pragma alloc_text(PAGE, CreateFilterDevice)
NTSTATUS CreateFilterDevice(const DEVICE_TYPE deviceType, PDEVICE_OBJECT* pNewFilterDevice)
{
    NTSTATUS ntStatus = STATUS_UNSUCCESSFUL;
    PFILE_HIDING_FILTER_DEVICE_EXTENSION pDeviceExt = NULL;

    PAGED_CODE();

    ntStatus = IoCreateDevice(GetFilterDriverObject(),
                              sizeof(FILE_HIDING_FILTER_DEVICE_EXTENSION),
                              NULL,
                              deviceType,
                              0,
                              FALSE,
                              pNewFilterDevice);

    if (NT_SUCCESS(ntStatus))
    {
        pDeviceExt = (PFILE_HIDING_FILTER_DEVICE_EXTENSION)(*pNewFilterDevice)->DeviceExtension;

        RtlZeroMemory(pDeviceExt, sizeof(FILE_HIDING_FILTER_DEVICE_EXTENSION));

        pDeviceExt->Dispatch.pDispatchClose = FilterDispatchClose;
        pDeviceExt->Dispatch.pDispatchCreate = FilterDispatchCreate;
        pDeviceExt->Dispatch.pDispatchDirectoryControl = FilterDispatchDirectoryControl;
    }

    return ntStatus;
}